import json
import os
from typing import Dict, Any, List, AsyncGenerator
import asyncio
from openai import AsyncOpenAI
from bson import ObjectId

class AIClient:
    """AI客户端基类 - 使用 OpenAI 官方库"""
    def __init__(self, api_key: str, base_url: str, default_model: str = "gpt-3.5-turbo"):
        """
        初始化 AI 客户端
        :param api_key: API 密钥
        :param base_url: API 基础 URL（例如：https://api.moonshot.cn/v1）
        :param default_model: 默认模型名称
        """
        # 确保 base_url 不包含 /chat/completions，OpenAI 库会自动添加
        if base_url.endswith('/chat/completions'):
            base_url = base_url.replace('/chat/completions', '')
        
        self.client = AsyncOpenAI(
            api_key=api_key,
            base_url=base_url
        )
        self.default_model = default_model
    
    async def chat_completion(self, messages: List[Dict[str, str]], model: str = None, **kwargs) -> Dict[str, Any]:
        """非流式聊天完成"""
        try:
            response = await self.client.chat.completions.create(
                model=model or self.default_model,
                messages=messages,
                **kwargs
            )
            return response.model_dump()
        except Exception as e:
            raise Exception(f"API请求失败: {str(e)}")
    
    async def chat_completion_stream(self, messages: List[Dict[str, str]], model: str = None, **kwargs) -> AsyncGenerator[str, None]:
        """流式聊天完成"""
        try:
            stream = await self.client.chat.completions.create(
                model=model or self.default_model,
                messages=messages,
                stream=True,
                **kwargs
            )
            
            async for chunk in stream:
                if chunk.choices and chunk.choices[0].delta.content:
                    yield chunk.choices[0].delta.content
                    
        except Exception as e:
            raise Exception(f"API流式请求失败: {str(e)}")


def get_ai_client(supplier_name: str, api_key: str, api_url: str) -> AIClient:
    """
    根据供应商名称获取对应的AI客户端
    
    :param supplier_name: 供应商名称（openai/kimi/moonshot/deepseek等）
    :param api_key: API 密钥
    :param api_url: API URL（可以是完整的或基础URL）
    :return: AI客户端实例
    """
    # 确保 api_url 是基础 URL（不包含 /chat/completions）
    if api_url.endswith('/chat/completions'):
        api_url = api_url.replace('/chat/completions', '')
    
    # 不同供应商的默认模型
    default_models = {
        "openai": "gpt-3.5-turbo",
        "kimi": "moonshot-v1-8k",
        "moonshot": "moonshot-v1-8k",
        "deepseek": "deepseek-chat",
        "anthropic": "claude-3-sonnet-20240229",
    }
    
    model = default_models.get(supplier_name.lower(), "gpt-3.5-turbo")
    return AIClient(api_key=api_key, base_url=api_url, default_model=model)